package org.jinghouyu.http.proxy.server;

import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.jinghouyu.http.proxy.stream.MergedInputStream;
import org.jinghouyu.http.proxy.utils.ParseUtils;

/**
 * response from server
 * @author liujingyu
 *
 */
public class Response implements Closeable {
	
	private LinkedHashMap<String, String> headers = new LinkedHashMap<>();
	private List<LinkedHashMap<String, String>> cookies = new ArrayList<>();
	private String messageHeader;
	private InputStream bodyInput;
	
	private boolean add(String headerStr) {
		int index = headerStr.indexOf(":");
		if(index <= 0) return false;
		String name = headerStr.substring(0, index).trim();
		String value = headerStr.substring(index + 1).trim();
		if("Set-Cookie".equals(name)) {
			parseCookies(value);
		} else {
			headers.put(name, value);
		}
		return true;
	}
	
	public String getHeader(String name) {
		if("Set-Cookie".equals(name)) return null;
		return headers.get(name);
	}
	
	public InputStream getBodyInput() {
		return bodyInput;
	}
	
	public int getStatus() {
		return Integer.parseInt(messageHeader.split(" ")[1]);
	}
	
	public String getProtocol() {
		return messageHeader.split(" ")[0];
	}
	
	private void parse_(InputStream in) throws IOException {
		byte[] firstLine = ParseUtils.readLine(in);
		messageHeader = new String(firstLine);
		parseHeaders(in);
		bodyInput = in;
	}
	
	private void parseHeaders(InputStream in) throws IOException {
		while(true) {
			byte[] line = ParseUtils.readLine(in);
			if(line.length == 0) {
				break;
			}
			add(new String(line));
		}
	}
	
	private void parseCookies(String cookieStr) {
		if(cookieStr == null) return;
		String[] cookiesString = cookieStr.split("; ");
		LinkedHashMap<String, String> map = new LinkedHashMap<>();
		for(String cookie : cookiesString) {
			int index = cookie.indexOf("=");
			if(index == -1) continue;
			String name = cookie.substring(0, index);
			String value = cookie.substring(index + 1);
			map.put(name, value);
		}
		cookies.add(map);
	}
	
	public static Response parse(InputStream in) throws IOException {
		Response response = new Response();
		response.parse_(in);
		return response;
	}
	
	private static final String CRLF = "\r\n";
	
	public String headersToString() {
		StringBuilder sb = new StringBuilder();
		for(String key : headers.keySet()) {
			String name = key;
			String value = headers.get(name);
			if(value == null) continue;
			sb.append(name)
			  .append(':')
			  .append(' ')
			  .append(value)
			  .append(CRLF);
		}
		sb.append(cookiesToString());
		return sb.toString();
	}
	
	private String cookiesToString() {
		if(cookies.size() == 0) return "";
		StringBuilder sb = new StringBuilder();
		int i = 0;
		for(Map<String, String> cookie : cookies) {
			if(cookie.isEmpty()) continue;
			StringBuilder cookieStr = new StringBuilder();
			cookieStr.append("Set-Cookie: ");
			for(String key : cookie.keySet()) {
				String value = cookie.get(key);
				cookieStr.append(key)
				  .append('=')
				  .append(value);
				if(i != cookie.size() - 1) {
					cookieStr.append(';')
					  .append(' ');
				}
				i++;
			}
			cookieStr.append(CRLF);
			sb.append(cookieStr);
		}
		return sb.toString();
	}
	
	/**
	 * to an inputstream
	 * @return
	 */
	public InputStream toHttpInputStream() {
		StringBuilder sb = new StringBuilder();
		sb.append(messageHeader)
		  .append(CRLF)
		  .append(headersToString())
		  .append(CRLF);
		return new MergedInputStream(bodyInput, sb.toString());
	}

	public void setHeaders(LinkedHashMap<String, String> headers) {
		this.headers = headers;
	}

	public void setMessageHeader(String messageHeader) {
		this.messageHeader = messageHeader;
	}

	public void setBodyInput(InputStream bodyInput) {
		this.bodyInput = bodyInput;
	}
	
	public void addHeader(String name, String value) {
		headers.put(name, value);
	}
	
	public void removeHeader(String name) {
		if("Set-Cookie".equals(name)) return;
		headers.remove(name);
	}
	
	public void addCookie(String name, String value) {
		boolean flag = false;
		for(Map<String, String> cookie : cookies) {
			if(cookie.containsKey(name)) {
				cookie.put(name, value);
				flag = true;
			}
		}
		if(!flag) {
			LinkedHashMap<String, String> cookie = null;
			if(cookies.size() == 0) {
				cookie = new LinkedHashMap<>();
				cookies.add(cookie);
			} else {
				cookie = cookies.get(cookies.size() - 1);
			}
			cookie.put(name, value);
		}
	}
	
	public void removeCookie(String name) {
		for(Map<String, String> cookie : cookies) {
			cookie.remove(name);
		}
	}
	
	public String getCookie(String name) {
		for(Map<String, String> cookie : cookies) {
			if(cookie.containsKey(name)) {
				return cookie.get(name);
			}
		}
		return null;
	}
	
	public Set<String> cookieKeySet() {
		Set<String> set = new LinkedHashSet<>();
		for(Map<String, String> cookie : cookies) {
			set.addAll(cookie.keySet());
		}
		return set;
	}
	
	public Set<String> headKeySet() {
		Set<String> set = headers.keySet();
		LinkedHashSet<String> keys = new LinkedHashSet<String>();
		for(String key : set) {
			if(!"Cookie".equals(key)) {
				keys.add(key);
			}
		}
		return keys;
	}

	@Override
	public void close() throws IOException {
		bodyInput.close();
	}
}